"""
一个鉴定nfsw内容的模块
"""
import keras
import numpy as np
from PIL import Image
from typing import Optional

from ._download import get_default_weights_path
from ._image import preprocess_image, Preprocessing
from ._model import make_open_nsfw_model

global_model: Optional[keras.models.Model] = None


def _update_global_model_if_needed(weights_path: Optional[str]) -> None:
    """
    地款模型
    :param weights_path:
    :return:
    """
    global global_model
    if global_model is None:
        global_model = make_open_nsfw_model(weights_path=weights_path)


def predict_image(
        image_path: str,
        preprocessing: Preprocessing = Preprocessing.YAHOO,
        weights_path: Optional[str] = get_default_weights_path(),
        grad_cam_path: Optional[str] = None,
        alpha: float = 0.8
) -> float:
    """
    Pipeline from single image path to predicted NSFW probability.
    Optionally generate and save the Grad-CAM plot.
    """
    pil_image = Image.open(image_path)
    # 图像处理
    image = preprocess_image(pil_image, preprocessing)
    _update_global_model_if_needed(weights_path)
    assert global_model is not None

    nsfw_probability = float(global_model(np.expand_dims(image, 0))[0][1])
    if grad_cam_path is not None:
        # TensorFlow will only be imported here.
        from ._inspection import make_and_save_nsfw_grad_cam

        make_and_save_nsfw_grad_cam(
            pil_image, preprocessing, global_model, grad_cam_path, alpha
        )

    return nsfw_probability
